'''
Created on Aug 14, 2009

@author: mkiyer
'''

import sys
import h5py
import numpy as np
from optparse import OptionParser
import os

def read_bed_file(fhd):
    for line in fhd:
        fields = line.strip().split('\t')
        yield (fields[0], int(fields[1]), int(fields[2]))

def plot_scatter_hist(x, y, outfile):
    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.ticker import NullFormatter
    
    nullfmt = NullFormatter()         # no labels
    # definitions for the axes
    left, width = 0.1, 0.65
    bottom, height = 0.1, 0.65
    bottom_h = left_h = left+width+0.02
    
    rect_scatter = [left, bottom, width, height]
    rect_histx = [left, bottom_h, width, 0.2]
    rect_histy = [left_h, bottom, 0.2, height]
    
    # start with a rectangular Figure
    plt.figure(1, figsize=(8,8))
    
    axScatter = plt.axes(rect_scatter)
    axHistx = plt.axes(rect_histx)
    axHisty = plt.axes(rect_histy)
    
    # no labels
    axHistx.xaxis.set_major_formatter(nullfmt)
    axHisty.yaxis.set_major_formatter(nullfmt)
    
    # the scatter plot:
    axScatter.scatter(x, y)
    
    # now determine nice limits by hand:
    binwidth = 0.25
    xymax = np.max( [np.max(np.fabs(x)), np.max(np.fabs(y))] )
    lim = ( int(xymax/binwidth) + 1) * binwidth
    
    axScatter.set_xlim( (-lim, lim) )
    axScatter.set_ylim( (-lim, lim) )
    
    bins = np.arange(-lim, lim + binwidth, binwidth)
    axHistx.hist(x, bins=bins)
    axHisty.hist(y, bins=bins, orientation='horizontal')
    
    axHistx.set_xlim( axScatter.get_xlim() )
    axHisty.set_ylim( axScatter.get_ylim() )    
    plt.savefig(outfile)
    plt.close()
    
def plot_scatter(xdata, ydata, outfile, 
                 x_name='Control Group',
                 y_name='Experimental Group',
                 title='Coverage scatter plot'):
    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plt.title(title)
    # size?
    # s = 20*np.log(np.array([x.ratio for x in intervals]))        
    # color?
    # c = [x.category * cm.jet.N / len(category_names)] * len(intervals)
    # c = x.category * cm.jet.N / len(category_names)
    plt.scatter(xdata, ydata, s=30, c='b', marker='o', cmap=None, norm=None,
                vmin=None, vmax=None, alpha=1.0, linewidths=None,
                verts=None)
    # label=category_names[category])
    #plt.legend()
    
    # best fit line
    x, y = sorted(xdata), sorted(ydata)    
    m = np.polyfit(x, y, 1)
    yfit = np.polyval(m, x)
    plt.plot(x, yfit, 'g')
    # correlation coefficient
    c = np.corrcoef(xdata, ydata)[0,1]
    r2 = c**2
    plt.figtext(0.815, 0.013, ' r^2=%.3f' % r2, color='black', weight='roman',
               size='small')    
    plt.grid()
    themax = max(x[-1], y[-1])
    plt.axis([-10, themax, -10, themax])
    plt.xlabel(x_name)
    plt.ylabel(y_name)
    plt.savefig(outfile)
    plt.close()

def plot_overlapping_hist(x, y, outfile):
    import matplotlib.pyplot as plt

    # histogram
    fig = plt.figure()

    # the histogram of the data
    n, bins, patches = plt.hist(x, facecolor='r')
    n2, bins2, patches2 = plt.hist(y, bins, facecolor='g', alpha=0.5)
    #n, bins, patches = plt.hist(exoncov, bins=50, log=True, facecolor='r', alpha=0.5)
    #n2, bins2, patches2 = plt.hist(covarray, bins, log=True, facecolor='g', alpha=0.5)
    #, normed=1, facecolor='g', alpha=0.75)
    #plt.legend([patches[0], patches2[0]], ['lander', 'all'])
    #n, bins, patches = plt.hist([lander_exoncov, refflat_exoncov, covarray], bins=50, range=(0, 50), log=True)

    #plt.xlabel('coverage')
    #plt.ylabel('log # of positions with said coverage')
    #plt.title('histogram of coverage')
    plt.grid(True)
    #plt.axis([40, 160, 0, 0.03])
    plt.savefig(outfile)
    plt.close()

def get_interval_coverage(intervals, h5file):
    h5data = h5py.File(h5file, 'r')
    for chrom, start, end in intervals:
        if chrom in h5data:
            dset = h5data[chrom]
            cov = 1000 * (np.sum(dset[start:end]) / float(end - start))
            yield cov
        else:
            yield 0.0    

if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] sample_file")
    optionparser.add_option("-x", dest="h5x",
                            help="h5 file containing experiment group data")
    optionparser.add_option("-y", dest="h5y",
                            help="h5 file containing control group data")
    optionparser.add_option("-o", "--output", dest="outfile",
                            default=None,
                            help="output file [default: %default]")
    (options, args) = optionparser.parse_args()
    if len(args) == 0:
        optionparser.error("no sample_file specified")
    interval_file = args[0]    

    if options.outfile == None:
        options.outfile = os.path.splitext(os.path.basename(interval_file))[0] + ".png"
    intervals = list(read_bed_file(open(interval_file)))

    # determine coverage of intervals
    h5data_x = h5py.File(options.h5x, 'r')
    h5data_y = h5py.File(options.h5y, 'r')
    
    xdata = []
    ydata = []
    i = 0
    m = 0
    for chrom, start, end in intervals:
        # count lines for status updates
        i+=1
        if i == 100000:
            m += 1
            sys.stderr.write(" %d\n" % (m*100000))
            i=0
        if chrom in h5data_x and chrom in h5data_y:
            dset_x = h5data_x[chrom]
            dset_y = h5data_y[chrom]        
            cov_x = 1000 * (np.sum(dset_x[start:end]) / float(end - start))
            cov_y = 1000 * (np.sum(dset_y[start:end]) / float(end - start))
            xdata.append(cov_x)
            ydata.append(cov_y)
        else:
            print '%s\t%d\t%d not found in hdf5 datasets, perhaps chrom name is bad' % (chrom, start, end)
            
    plot_scatter(xdata, ydata, os.path.basename(options.outfile), 
                 x_name=os.path.basename(options.h5x),
                 y_name=os.path.basename(options.h5y),
                 title="Coverage for intervals in %s" % os.path.basename(interval_file))
    
    #plot_scatter_hist(xdata, ydata, "scatterhist.png")                     
#    ratio = float(cov_t) / float(cov_c)
#    options.outfile.write("%s\t%d\t%d\t%.2f\t%.2f\t%.2f\n" %
#                          (chrom, start, end,
#                           cov_t, cov_c, ratio))
